#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
树模块

此模块提供用于表示RRT树结构的Tree类。
"""

import numpy as np
from .node import Node


class Tree:
    """
    表示RRT树结构的类。
    """
    
    def __init__(self, root_node):
        """
        使用根节点初始化树。
        
        参数:
            root_node (Node): 树的根节点
        """
        self.root = root_node
        self.nodes = [root_node]
        self.size = 1
    
    def add_node(self, node):
        """
        向树中添加节点。
        
        参数:
            node (Node): 要添加到树中的节点
        """
        self.nodes.append(node)
        self.size += 1
        
        # 将该节点作为子节点添加到其父节点
        if node.parent:
            node.parent.add_child(node)
    
    def get_nearest_node(self, x, y):
        """
        查找树中距给定坐标最近的节点。
        
        参数:
            x (float): X坐标
            y (float): Y坐标
            
        返回:
            Node: 树中最近的节点
        """
        # 创建临时节点用于距离计算
        target_node = Node(x, y)
        
        # 查找具有最小距离的节点
        min_distance = float('inf')
        nearest_node = None
        
        for node in self.nodes:
            distance = node.distance_to(target_node)
            if distance < min_distance:
                min_distance = distance
                nearest_node = node
                
        return nearest_node
    
    def get_nodes_in_radius(self, x, y, radius):
        """
        查找指定点给定半径内的所有节点。
        
        参数:
            x (float): 中心点的X坐标
            y (float): 中心点的Y坐标
            radius (float): 搜索半径
            
        返回:
            list: 半径内的节点列表
        """
        # 创建临时节点用于距离计算
        center_node = Node(x, y)
        
        # 查找半径内的所有节点
        nodes_in_radius = []
        
        for node in self.nodes:
            if node.distance_to(center_node) <= radius:
                nodes_in_radius.append(node)
                
        return nodes_in_radius
    
    def get_path_to_node(self, node):
        """
        获取从根节点到指定节点的路径。
        
        参数:
            node (Node): 目标节点
            
        返回:
            list: 从根节点到目标节点的(x, y)坐标列表
        """
        path = []
        current = node
        
        # 从节点向上遍历到根节点
        while current:
            path.append(current.position)
            current = current.parent
            
        # 反转路径以获得从根节点到目标节点的顺序
        return path[::-1]
    
    def get_path_nodes(self, node):
        """
        获取从根节点到指定节点的路径节点。
        
        参数:
            node (Node): 目标节点
            
        返回:
            list: 从根节点到目标节点的Node对象列表
        """
        path_nodes = []
        current = node
        
        # 从节点向上遍历到根节点
        while current:
            path_nodes.append(current)
            current = current.parent
            
        # 反转路径以获得从根节点到目标节点的顺序
        return path_nodes[::-1]
    
    def compute_path_length(self, node):
        """
        计算从根节点到指定节点的路径长度。
        
        参数:
            node (Node): 目标节点
            
        返回:
            float: 路径长度（欧几里得距离之和）
        """
        if node is None:
            return float('inf')
            
        return node.cost
    
    def get_node_positions(self):
        """
        获取树中所有节点的位置。
        
        返回:
            tuple: (xs, ys) x和y坐标的列表
        """
        xs = [node.x for node in self.nodes]
        ys = [node.y for node in self.nodes]
        return xs, ys
    
    def get_edges(self):
        """
        获取树中所有边用于可视化。
        
        返回:
            list: 表示边的((x1, y1), (x2, y2))对的列表
        """
        edges = []
        
        # 跳过没有父节点的根节点
        for node in self.nodes[1:]:
            if node.parent:
                edges.append((
                    (node.parent.x, node.parent.y),
                    (node.x, node.y)
                ))
                
        return edges 